{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import sklearn\n",
    "import numpy as np\n",
    "import os,json\n",
    "\n",
    "from alphagen.data.expression import *\n",
    "from alphagen.models.alpha_pool import AlphaPool\n",
    "from alphagen.utils import reseed_everything\n",
    "from alphagen_generic.operators import funcs as generic_funcs\n",
    "from alphagen_generic.features import *\n",
    "from gan.utils.data import get_data_by_year"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "def get_ml_data(data):\n",
    "    df = data.df_bak.copy()\n",
    "    df.columns = ['open','close','high','low','volume','vwap']\n",
    "    vwap_unstack = df['vwap'].unstack()\n",
    "    tmp = (vwap_unstack.shift(-21)/vwap_unstack.shift(-1))-1\n",
    "    label = tmp.stack().reindex(df.index)\n",
    "    df['label'] = label\n",
    "\n",
    "    feature = df[['open','close','high','low','volume','vwap']]\n",
    "    tmp = feature.unstack()\n",
    "    feature = (tmp/tmp.shift(1)-1)#.stack().reindex(df.index)\n",
    "\n",
    "    result_feature = []\n",
    "    cur = feature.stack().reindex(df.index)\n",
    "    cur.columns = [f'{col}0' for col in cur.columns]\n",
    "    for past in tqdm(range(1,60)):\n",
    "        cur = feature.shift(past).stack().reindex(df.index)\n",
    "        cur.columns = [f'{col}{past}' for col in cur.columns]\n",
    "        result_feature.append(cur)\n",
    "    result_feature = pd.concat(result_feature,axis=1)\n",
    "    df = pd.concat([result_feature,df['label']],axis=1)\n",
    "    start_date = data._dates[data.max_backtrack_days]\n",
    "    end_date = data._dates[-data.max_future_days]\n",
    "    return df.loc[start_date:end_date]\n",
    "\n",
    "def normalize_data(df_train, df_valid, df_test):\n",
    "    # Get the column names of the features\n",
    "    labels = [df_train.iloc[:, -1], df_valid.iloc[:, -1], df_test.iloc[:, -1]]\n",
    "    df_train_features = df_train.iloc[:, :-1]\n",
    "    df_valid_features = df_valid.iloc[:, :-1]\n",
    "    df_test_features = df_test.iloc[:, :-1]\n",
    "\n",
    "    _mean = df_train_features.mean()\n",
    "    _std = df_train_features.std()\n",
    "    print('1')\n",
    "    df_train_norm = (df_train_features - _mean) / _std\n",
    "    print('2')\n",
    "    df_valid_norm = (df_valid_features - _mean) / _std\n",
    "    print('3')\n",
    "    df_test_norm = (df_test_features - _mean) / _std\n",
    "\n",
    "    df_train_norm.fillna(0, inplace=True)\n",
    "    df_valid_norm.fillna(0, inplace=True)\n",
    "    df_test_norm.fillna(0, inplace=True)\n",
    "\n",
    "\n",
    "    df_train_norm['label'] = np.nan_to_num(labels[0],nan=0,posinf=0,neginf=0)\n",
    "    df_valid_norm['label'] = np.nan_to_num(labels[1],nan=0,posinf=0,neginf=0)\n",
    "    df_test_norm['label'] = np.nan_to_num(labels[2],nan=0,posinf=0,neginf=0)\n",
    "\n",
    "    df_train_norm['label'] = df_train_norm['label'].groupby('datetime').transform(lambda x: (x - x.mean()) / x.std()).clip(-4, 4)\n",
    "\n",
    "    df_train_norm = df_train_norm.clip(-4, 4)\n",
    "    df_valid_norm = df_valid_norm.clip(-4, 4)\n",
    "    df_test_norm = df_test_norm.clip(-4, 4)\n",
    "\n",
    "    return df_train_norm, df_valid_norm, df_test_norm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Lightgbm Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "import lightgbm as lgb\n",
    "import pandas as pd\n",
    "\n",
    "def train_lightgbm_model(df_train, df_valid, df_test):\n",
    "    # Fill NaN values with 0\n",
    "    df_train_filled = df_train.fillna(0)\n",
    "    df_valid_filled = df_valid.fillna(0)\n",
    "    df_test_filled = df_test.fillna(0)\n",
    "\n",
    "    # Separate features and labels\n",
    "    X_train = df_train_filled.drop(columns=['label'])\n",
    "    y_train = df_train_filled['label']\n",
    "    X_valid = df_valid_filled.drop(columns=['label'])\n",
    "    y_valid = df_valid_filled['label']\n",
    "    X_test = df_test_filled.drop(columns=['label'])\n",
    "    y_test = df_test_filled['label']\n",
    "\n",
    "    # Convert data to LightGBM Dataset format\n",
    "    train_data = lgb.Dataset(X_train, label=y_train)\n",
    "    valid_data = lgb.Dataset(X_valid, label=y_valid)\n",
    "\n",
    "    # Set hyperparameters for LightGBM model\n",
    "    params = {\n",
    "        'objective': 'regression',\n",
    "        'metric': 'mse',\n",
    "        'num_leaves': 210,\n",
    "        'max_depth': 8,\n",
    "        'learning_rate': 0.05,\n",
    "        'feature_fraction': 0.9,\n",
    "        'bagging_fraction': 0.8,\n",
    "        'bagging_freq': 5,\n",
    "        'verbose': 0\n",
    "    }\n",
    "\n",
    "    # Train the LightGBM model\n",
    "    model = lgb.train(params, train_data, valid_sets=[train_data, valid_data], num_boost_round=1000, early_stopping_rounds=100, verbose_eval=100)\n",
    "\n",
    "    # Evaluate the model on the test set\n",
    "    y_pred = model.predict(X_test)\n",
    "    \n",
    "    pred = pd.concat([pd.Series(y_pred,index=df_test.index),df_test['label']],axis=1)\n",
    "    # Print the RMSE score\n",
    "    # rmse = np.sqrt(mean_squared_error(y_test, y_pred))\n",
    "    # print(f\"RMSE: {rmse}\")\n",
    "\n",
    "    return  model,pred\n",
    "\n",
    "for instruments in ['csi300','csi500']:\n",
    "    for train_end in range(2016,2021):\n",
    "        returned = get_data_by_year(\n",
    "            train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2,\n",
    "            instruments=instruments, target=target,freq='day',)\n",
    "        data_all, data,data_valid,data_valid_withhead,data_test,data_test_withhead,_ = returned\n",
    "        df_train = get_ml_data(data)\n",
    "        df_valid = get_ml_data(data_valid)\n",
    "        df_test = get_ml_data(data_test)\n",
    "        df_train, df_valid, df_test = normalize_data(df_train, df_valid, df_test)\n",
    "        \n",
    "        model_name = 'lgbm'\n",
    "        name = f\"{instruments}_{model_name}_{train_end}\"\n",
    "        os.makedirs(f\"out_ml/{name}\",exist_ok=True)\n",
    "        model,pred = train_lightgbm_model(df_train, df_valid, df_test)\n",
    "        model.save_model(f\"out_ml/{name}/{model_name}.pt\")\n",
    "        pred.to_pickle(f\"out_ml/{name}/pred.pkl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Xgboost Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import xgboost as xgb\n",
    "import pandas as pd\n",
    "\n",
    "def train_xgboost_model(df_train, df_valid, df_test):\n",
    "    # Fill NaN values with 0\n",
    "    df_train_filled = df_train.fillna(0)\n",
    "    df_valid_filled = df_valid.fillna(0)\n",
    "    df_test_filled = df_test.fillna(0)\n",
    "\n",
    "    # Separate features and labels\n",
    "    X_train = df_train_filled.drop(columns=['label'])\n",
    "    y_train = df_train_filled['label']\n",
    "    X_valid = df_valid_filled.drop(columns=['label'])\n",
    "    y_valid = df_valid_filled['label']\n",
    "    X_test = df_test_filled.drop(columns=['label'])\n",
    "    y_test = df_test_filled['label']\n",
    "\n",
    "    # Convert data to DMatrix format\n",
    "    dtrain = xgb.DMatrix(X_train, label=y_train)\n",
    "    dvalid = xgb.DMatrix(X_valid, label=y_valid)\n",
    "\n",
    "    # Set hyperparameters for XGBoost model\n",
    "    params = {\n",
    "        'objective': 'reg:squarederror',\n",
    "        'eval_metric': 'rmse',\n",
    "        'colsample_bytree': 0.8879,\n",
    "        'learning_rate': 0.2,\n",
    "        'subsample': 0.8789,\n",
    "        'lambda': 205.6999,\n",
    "        'alpha': 580.9768,\n",
    "        'max_depth': 8,\n",
    "        'num_boost_round': 1000,\n",
    "        'early_stopping_rounds': 100,\n",
    "        'verbose_eval': 100\n",
    "    }\n",
    "\n",
    "    # Train the XGBoost model\n",
    "    model = xgb.train(params, dtrain, evals=[(dtrain, 'train'), (dvalid, 'valid')], early_stopping_rounds=params['early_stopping_rounds'], verbose_eval=params['verbose_eval'])\n",
    "\n",
    "    # Convert test data to DMatrix format\n",
    "    dtest = xgb.DMatrix(X_test)\n",
    "\n",
    "    # Make predictions on the test set\n",
    "    y_pred = model.predict(dtest)\n",
    "\n",
    "    # Combine the predictions with the actual labels\n",
    "    # pred = pd.concat([df_test['label'], pd.Series(y_pred, index=df_test.index)], axis=1)\n",
    "    pred = pd.concat([pd.Series(y_pred,index=df_test.index),df_test['label']],axis=1)\n",
    "\n",
    "    return model, pred\n",
    "\n",
    "\n",
    "\n",
    "for instruments in ['csi300','csi500']:\n",
    "    for train_end in range(2016,2021):\n",
    "        returned = get_data_by_year(\n",
    "            train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2,\n",
    "            instruments=instruments, target=target,freq='day',)\n",
    "        data_all, data,data_valid,data_valid_withhead,data_test,data_test_withhead,_ = returned\n",
    "        df_train = get_ml_data(data)\n",
    "        df_valid = get_ml_data(data_valid)\n",
    "        df_test = get_ml_data(data_test)\n",
    "        df_train, df_valid, df_test = normalize_data(df_train, df_valid, df_test)\n",
    "        \n",
    "        model_name = 'xgb'\n",
    "        name = f\"{instruments}_{model_name}_{train_end}\"\n",
    "        os.makedirs(f\"out_ml/{name}\",exist_ok=True)\n",
    "        model,pred = train_xgboost_model(df_train, df_valid, df_test)\n",
    "        model.save_model(f\"out_ml/{name}/{model_name}.pt\")\n",
    "        pred.to_pickle(f\"out_ml/{name}/pred.pkl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train MLP Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from sklearn.metrics import mean_squared_error\n",
    "\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "def train_mlp_model(df_train, df_valid, df_test):\n",
    "    # Fill NaN values with 0\n",
    "    df_train_filled = df_train.fillna(0)\n",
    "    df_valid_filled = df_valid.fillna(0)\n",
    "    df_test_filled = df_test.fillna(0)\n",
    "\n",
    "    # Separate features and labels\n",
    "    X_train = df_train_filled.drop(columns=['label']).values\n",
    "    y_train = df_train_filled['label'].values\n",
    "    X_valid = df_valid_filled.drop(columns=['label']).values\n",
    "    y_valid = df_valid_filled['label'].values\n",
    "    X_test = df_test_filled.drop(columns=['label']).values\n",
    "    y_test = df_test_filled['label'].values\n",
    "\n",
    "    # Convert data to tensors\n",
    "    X_train = torch.tensor(X_train, dtype=torch.float32)\n",
    "    y_train = torch.tensor(y_train, dtype=torch.float32)\n",
    "    X_valid = torch.tensor(X_valid, dtype=torch.float32)\n",
    "    y_valid = torch.tensor(y_valid, dtype=torch.float32)\n",
    "    X_test = torch.tensor(X_test, dtype=torch.float32)\n",
    "    y_test = torch.tensor(y_test, dtype=torch.float32)\n",
    "\n",
    "    # Define the MLP model\n",
    "    model = nn.Sequential(\n",
    "        nn.Linear(X_train.shape[1], 64),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(64, 32),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(32, 1)\n",
    "    )\n",
    "\n",
    "    # Define the loss function and optimizer\n",
    "    criterion = nn.MSELoss()\n",
    "    optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "\n",
    "    # Move the model to GPU if available\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    model.to(device)\n",
    "\n",
    "    # Train the model\n",
    "    num_epochs = 10\n",
    "    batch_size = 512\n",
    "    for epoch in range(num_epochs):\n",
    "        # Shuffle the training data\n",
    "        indices = torch.randperm(X_train.shape[0])\n",
    "        X_train_shuffled = X_train[indices]\n",
    "        y_train_shuffled = y_train[indices]\n",
    "\n",
    "        # Mini-batch training\n",
    "        for i in tqdm(range(0, X_train.shape[0], batch_size)):\n",
    "            # Get the mini-batch\n",
    "            X_batch = X_train_shuffled[i:i+batch_size]\n",
    "            y_batch = y_train_shuffled[i:i+batch_size]\n",
    "\n",
    "            # Move the mini-batch to GPU if available\n",
    "            X_batch = X_batch.to(device)\n",
    "            y_batch = y_batch.to(device)\n",
    "\n",
    "            # Forward pass\n",
    "            outputs = model(X_batch)\n",
    "            loss = criterion(outputs.flatten(), y_batch.flatten())\n",
    "\n",
    "            # Backward pass and optimization\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "    # Evaluate the model on the test set\n",
    "    with torch.no_grad():\n",
    "        # Move the test data to GPU if available\n",
    "        test_outputs = model(X_test.to(device)).detach().cpu().numpy().flatten()\n",
    "        pred_df = pd.concat([df_test['label'],pd.Series(test_outputs,index=df_test.index)],axis=1)\n",
    "    torch.cuda.empty_cache()\n",
    "    return model, pred_df\n",
    "\n",
    "for instruments in ['csi300','csi500']:\n",
    "    for train_end in range(2016,2021):\n",
    "        returned = get_data_by_year(\n",
    "            train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2,\n",
    "            instruments=instruments, target=target,freq='day',)\n",
    "        data_all, data,data_valid,data_valid_withhead,data_test,data_test_withhead,_ = returned\n",
    "        df_train = get_ml_data(data)\n",
    "        df_valid = get_ml_data(data_valid)\n",
    "        df_test = get_ml_data(data_test)\n",
    "        df_train, df_valid, df_test = normalize_data(df_train, df_valid, df_test)\n",
    "        for seed in range(5):\n",
    "            reseed_everything(seed)\n",
    "            model_name = 'mlp'\n",
    "            name = f\"{instruments}_{model_name}_{train_end}\"\n",
    "            os.makedirs(f\"out_ml/{name}\",exist_ok=True)\n",
    "            model,pred = train_mlp_model(df_train, df_valid, df_test)\n",
    "            # model.save_model(f\"out_ml/{name}/model.pt\")\n",
    "            pred.to_pickle(f\"out_ml/{name}/pred_{seed}.pkl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Show LightGbm Result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "instruments = 'csi300'\n",
    "\n",
    "result = []\n",
    "for year in range(2016,2021):\n",
    "    result.append(pd.read_pickle(f'out_ml/{instruments}_lgb_{year}/pred.pkl'))\n",
    "df = pd.concat(result,0)\n",
    "\n",
    "def get_final_metrics(df):\n",
    "    ic_raw = df.groupby('datetime').corr()['label'].unstack().iloc[:,:-1]\n",
    "    ic_s = ic_raw.mean(axis=0)\n",
    "    ic_s_std = ic_raw.std(axis=0)\n",
    "    icir_s = ic_s/ic_s_std\n",
    "\n",
    "    ric_raw = df.groupby('datetime').corr('spearman')['label'].unstack().iloc[:,:-1]\n",
    "    ric_s = ric_raw.mean(axis=0)\n",
    "    ric_s_std = ric_raw.std(axis=0)\n",
    "    ricir_s = ric_s/ric_s_std\n",
    "    return {'IC':ic_s.mean(),'ICIR':icir_s.mean(),'RankIC':ric_s.mean(),'RankICIR':ricir_s.mean()}\n",
    "\n",
    "print('LightGBM Result:\\n')\n",
    "print(get_final_metrics(df))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Show XGBoost Result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "instruments = 'csi300'\n",
    "\n",
    "result = []\n",
    "for year in range(2016,2021):\n",
    "    result.append(pd.read_pickle(f'out_ml/{instruments}_xgb_{year}/pred.pkl'))\n",
    "df = pd.concat(result,0)\n",
    "\n",
    "def get_final_metrics(df):\n",
    "    ic_raw = df.groupby('datetime').corr()['label'].unstack().iloc[:,:-1]\n",
    "    ic_s = ic_raw.mean(axis=0)\n",
    "    ic_s_std = ic_raw.std(axis=0)\n",
    "    icir_s = ic_s/ic_s_std\n",
    "\n",
    "    ric_raw = df.groupby('datetime').corr('spearman')['label'].unstack().iloc[:,:-1]\n",
    "    ric_s = ric_raw.mean(axis=0)\n",
    "    ric_s_std = ric_raw.std(axis=0)\n",
    "    ricir_s = ric_s/ric_s_std\n",
    "    return {'IC':ic_s.mean(),'ICIR':icir_s.mean(),'RankIC':ric_s.mean(),'RankICIR':ricir_s.mean()}\n",
    "\n",
    "print('XGBoost Result:\\n')\n",
    "print(get_final_metrics(df))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Show MLP Result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "result_all = []\n",
    "for seed in range(5):\n",
    "    result = []\n",
    "    for year in range(2016,2021):\n",
    "        result.append(pd.read_pickle(f'out_ml/csi500_mlp_{year}/pred_{seed}.pkl'))\n",
    "    df = pd.concat(result,0)#.groupby('datetime').corr('spearman')['label'].unstack().mean()\n",
    "    df.columns = ['pred','label']\n",
    "    df=df[['pred','label']]\n",
    "\n",
    "\n",
    "\n",
    "    def get_final_metrics(df):\n",
    "        ic_raw = df.groupby('datetime').corr()['label'].unstack().iloc[:,:-1]\n",
    "        ic_s = ic_raw.mean(axis=0)\n",
    "        ic_s_std = ic_raw.std(axis=0)\n",
    "        icir_s = ic_s/ic_s_std\n",
    "\n",
    "        ric_raw = df.groupby('datetime').corr('spearman')['label'].unstack().iloc[:,:-1]\n",
    "        ric_s = ric_raw.mean(axis=0)\n",
    "        ric_s_std = ric_raw.std(axis=0)\n",
    "        ricir_s = ric_s/ric_s_std\n",
    "        return {'IC':ic_s.mean(),'ICIR':icir_s.mean(),'RankIC':ric_s.mean(),'RankICIR':ricir_s.mean()}\n",
    "    tmp = get_final_metrics(df)\n",
    "    result_all.append(tmp)\n",
    "print('MLP Result:\\n')\n",
    "print(result_all)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
